/* Copyright (c) 2008, RPI-ACM, Jesse Kempf <jessekempf@gmail.com> 
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 * - Redistributions of source code must retain the above copyright notice,
 *   this list of conditions and the following disclaimer.
 * - Redistributions in binary form must reproduce the above copyright notice,
 *   this list of conditions and the following disclaimer in the documentation
 *   and/or other materials provided with the distribution.
 * - Neither the name of RPI-ACM nor the names of its contributors may
 *   be used to endorse or promote products derived from this software without
 *   specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <Python.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/uio.h>

const uint8_t NONCE = 0xAB;

static PyObject *passfd_error;

PyDoc_STRVAR(module_doc, "PyNIX File Descriptor Passing Module.\n"
"\n"
"Unix systems provide the ability to pass file descriptors over\n"
"domain sockets. The passfd module provides utility functions\n"
"to make file descriptor passing trivial.\n");

PyDoc_STRVAR(passfd_doc, "passfd(sockno, fd)\n"
"\n"
"Send the given file descriptor across the indicated socket.");

static PyObject *passfd(PyObject *self, PyObject *args) {
	struct iovec iov[1];
	struct msghdr *msg = PyMem_Malloc(sizeof(*msg));
	struct cmsghdr *cmsg;
	int sock, fd, rv;

	if (msg == NULL)
		return PyErr_NoMemory();

	memset(msg, 0, sizeof(*msg));

	/*
	 * Unbox parameters.
	 */

	if (!PyArg_ParseTuple(args, "ii", &sock, &fd))
		return (NULL);

	/*
	* Set up the message.
	*/

	msg->msg_name = NULL;
	msg->msg_namelen = 0;

	msg->msg_controllen = CMSG_SPACE(sizeof(int));
	msg->msg_control = PyMem_Malloc(msg->msg_controllen);
	if (msg->msg_control == NULL)
		return PyErr_NoMemory();

	/* Set the one-octet nonce to be passed. */
	iov[0].iov_base = (void *)&NONCE;
	iov[0].iov_len = sizeof(NONCE);

	msg->msg_iov = iov;
	msg->msg_iovlen = 1;

	/* Set the control message particulars. */
	cmsg = CMSG_FIRSTHDR(msg);
	cmsg->cmsg_len = CMSG_LEN(sizeof(int));
	cmsg->cmsg_level = SOL_SOCKET;
	cmsg->cmsg_type = SCM_RIGHTS;
	(*(int *) CMSG_DATA(cmsg)) = fd;

	/* Send the fd across the socket. */
	rv = sendmsg(sock, msg, 0);

	PyMem_Free(msg->msg_control);
	PyMem_Free(msg);

	if (rv == -1) {
		PyErr_SetFromErrno(PyExc_OSError);
		return (NULL);
	}

	Py_RETURN_NONE;
}

PyDoc_STRVAR(grabfd_doc, "grabfd(sockno) -> fd\n"
"\n"
"Receive the file descriptor sent across the indicated socket.");

static PyObject *grabfd(PyObject *self, PyObject *args) {
	struct iovec iov[1];
	struct msghdr *msg = PyMem_Malloc(sizeof(*msg));
	struct cmsghdr *cmsg;
	int sock, fd, rv;
	uint8_t nonce;

	if (msg == NULL)
		return PyErr_NoMemory();
	
	memset(msg, 0, sizeof(*msg));

	/*
	 * Unbox parameters.
	 */

	if (!PyArg_ParseTuple(args, "i", &sock))
		return (NULL);

	msg->msg_controllen = CMSG_SPACE(sizeof(int));
	msg->msg_control = PyMem_Malloc(msg->msg_controllen);
	if (msg->msg_control == NULL)
		return PyErr_NoMemory();

	msg->msg_name = NULL;
	msg->msg_namelen = 0;

	iov[0].iov_base = &nonce;
	iov[0].iov_len = sizeof(nonce);

	msg->msg_iov = iov;
	msg->msg_iovlen = 1;

	rv = recvmsg(sock, msg, 0);

	if (rv == -1) {
		PyErr_SetFromErrno(PyExc_OSError);
		return (NULL);
	}

	if (nonce != NONCE) {
		PyErr_SetString(passfd_error, "incorrect nonce.");
		return (NULL);
	}

	cmsg = CMSG_FIRSTHDR(msg);

	if (cmsg == NULL || cmsg->cmsg_len != CMSG_LEN(sizeof(int))
		|| cmsg->cmsg_level != SOL_SOCKET || cmsg->cmsg_type != SCM_RIGHTS) {
			PyErr_SetString(passfd_error, "incorrect message type received.");
			return (NULL);
	}

	fd = *((int *) CMSG_DATA(cmsg));

	PyMem_Free(msg->msg_control);
	PyMem_Free(msg);

	return (Py_BuildValue("i", fd));

}

static PyMethodDef passfdMethods[] = {
	{"passfd", passfd, METH_VARARGS, passfd_doc},
	{"grabfd", grabfd, METH_VARARGS, grabfd_doc},
	{NULL, NULL, 0, NULL}
};

PyMODINIT_FUNC initpassfd(void) {
	PyObject *m;

	m = Py_InitModule3("passfd", passfdMethods, module_doc);
	passfd_error = PyErr_NewException("passfd.error", NULL, NULL);
	if (passfd_error == NULL)
		return;
	Py_INCREF(passfd_error);
	PyModule_AddObject(m, "error", passfd_error);
}
